package com.geo.geotool.service.grid;

import com.geo.geotool.contour.algorithm.Kriging;
import lombok.extern.slf4j.Slf4j;
import org.geotools.coverage.grid.GridCoverage2D;
import org.geotools.coverage.grid.GridCoverageFactory;
import org.geotools.data.simple.SimpleFeatureCollection;
import org.geotools.geometry.Envelope2D;
import org.geotools.process.raster.CropCoverage;
import org.geotools.process.raster.PolygonExtractionProcess;
import org.geotools.process.raster.ScaleCoverage;
import org.jaitools.numeric.Range;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.springframework.stereotype.Component;

import javax.media.jai.RasterFactory;
import java.awt.image.DataBuffer;
import java.awt.image.WritableRaster;
import java.io.IOException;
import java.util.List;

/**
 * @Description: 栅格服务
 * @Author: 逃离地球
 */

@Component
@Slf4j
public class GridService {
    /**
     *
     * @param data  double[point.length][3]二维数组 --> point[0]为x坐标，point[1]为y坐标，point[2]为插值的值
     * @param height 栅格的高度
     * @param width 栅格的宽度
     * @param xMin 栅格的x最小值
     * @param xMax 栅格的x最大值
     * @param yMin 栅格的y最小值
     * @param yMax 栅格的y最大值
     * @param crs 坐标系
     * @return
     */
    public GridCoverage2D kriging(double[][] data, int height, int width , double xMin, double xMax, double yMin, double yMax, CoordinateReferenceSystem crs) {
        // 构建克里金插值分析
        Kriging kriging = new Kriging(Kriging.SPHERICAL_MODEL, 0, 100);
        // 获取需要插值的属性值数组
        double[] targetValues = new double[data.length];
        double[] xList = new double[data.length];
        double[] yList = new double[data.length];
        for(int i = 0; i < data.length; i++) {
            double[] point = data[i];
            targetValues[i] = point[2];
            xList[i] = point[0];
            yList[i] = point[1];
        }
        // 训练模型
        kriging.train(targetValues, xList, yList);
        // 创建栅格
        WritableRaster raster = RasterFactory.
                createBandedRaster(DataBuffer.TYPE_DOUBLE, width, height, 1, null);
        // 将数据写入 WritableRaster 对象中
        for (int x = 0; x < width; x++) {
            for (int y = 0; y < height; y++) {
                double zbx = xMin + (xMax - xMin) * x / width;
                //raster 是从左上角开始的，所以这里从上往下
                double zby = yMax - (yMax - yMin) * y / height;
                double value  = kriging.predict(zbx,zby);
                raster.setSample(x,  y, 0, value);
            }
        }
        // 构建gridCoverage2d
        GridCoverageFactory factory = new GridCoverageFactory();
        GridCoverage2D coverage = factory.create("contour_polygon", raster,
                new Envelope2D(crs, xMin, yMin, xMax - xMin, yMax - yMin));
        return coverage;
    }


    /**
     * 裁剪栅格
     * @param coverage 源栅格
     * @param geometry 裁剪区域
     * @return
     * @throws IOException
     */
    public GridCoverage2D crop(GridCoverage2D coverage,  Geometry  geometry) throws IOException {
        CropCoverage cropCoverage = new CropCoverage();
        // 裁剪,返回裁剪后的栅格
        // progressListener 为进度监听器，可以为null`
        GridCoverage2D crop = cropCoverage.execute(coverage,geometry,null);
        return crop;
    }

    /**
     * 缩放重采样
     * @return
     */
    public GridCoverage2D scale(GridCoverage2D coverage,double xScale,double yScale)  throws Exception{
        ScaleCoverage scaleCoverage = new ScaleCoverage();
        // 不设置插值方式，默认为最近邻插值
        GridCoverage2D scale = scaleCoverage.execute(coverage,xScale,yScale,0,0,null);
        return scale;
    }


    /**
     * 栅格转等值面
     * @param coverage
     * @param ranges
     * @return
     * @throws Exception
     */
    public SimpleFeatureCollection polygonExtract(GridCoverage2D coverage,List<Range> ranges)  throws Exception{
        PolygonExtractionProcess polygonExtractionProcess = new PolygonExtractionProcess();
        SimpleFeatureCollection polygonFeatures = polygonExtractionProcess.execute(coverage,0,true,null,null,ranges,null);
        return polygonFeatures;
    }
}
